/******************************************************************/
/*                                                                */
/*  Winpooch : Windows Watchdog                                   */
/*  Copyright (C) 2004-2005  Benoit Blanchon                      */
/*                                                                */
/*  This program is free software; you can redistribute it        */
/*  and/or modify it under the terms of the GNU General Public    */
/*  License as published by the Free Software Foundation; either  */
/*  version 2 of the License, or (at your option) any later       */
/*  version.                                                      */
/*                                                                */
/*  This program is distributed in the hope that it will be       */
/*  useful, but WITHOUT ANY WARRANTY; without even the implied    */
/*  warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR       */
/*  PURPOSE.  See the GNU General Public License for more         */
/*  details.                                                      */
/*                                                                */
/*  You should have received a copy of the GNU General Public     */
/*  License along with this program; if not, write to the Free    */
/*  Software Foundation, Inc.,                                    */
/*  675 Mass Ave, Cambridge, MA 02139, USA.                       */
/*                                                                */
/******************************************************************/


/******************************************************************/
/* Build configuration                                            */
/******************************************************************/

#define SPYDLL_BUILD_DLL

#define	TRACE_LEVEL	2

//#define SAVE_BOOT_INFO

// adds a break point into boot code used when injecting suspended threads
//#define BREAK_IN_BOOT


/******************************************************************/
/* Includes                                                       */
/******************************************************************/

// module's interface
#include "SpyDll.h"

// standard headers
#include <windows.h>
#include <stdio.h>
#include <imagehlp.h>
#include <tchar.h>
#include <psapi.h>

// project's headers
#include "Assert.h"
#include "Hooks.h"
#include "HookReg.h"
#include "SpyDllInt.h"
#include "Trace.h"

// TO BE REMOVED
#include "HookSys.h"


/******************************************************************/
/* Internal macros                                                */
/******************************************************************/

#define SHARED __attribute__((section(".shr"),shared))

#define arraysize(A) (sizeof(A)/sizeof((A)[0]))


/******************************************************************/
/* Internal constants                                             */
/******************************************************************/

#define LOADLIBRARYW_BAD_ADDRESS	0x7788786F

#ifdef UNICODE
  #define LOADLIBRARY_ANSI "LoadLibraryW"
  #define LOADLIBRARYEX_ANSI "LoadLibraryExW"
#else
  #define LOADLIBRARY_ANSI "LoadLibraryA"
  #define LOADLIBRARYEX_ANSI "LoadLibraryExA"
#endif


/******************************************************************/
/* Internal data                                                  */
/******************************************************************/

DWORD	sg_dwServerProcessID	SHARED = 0  ;
PROC	sg_pfnNotifyProc	SHARED = NULL ;

HINSTANCE	g_hInstance = NULL ;  // instance of this dll
HMODULE		g_hModule = NULL ; // module where this dll is attached
MODULEINFO	g_miInstance ;

// i made this global because GetProcAddress return bad values 
// when injecting suspended processes.
PROC		g_pfnLoadLibrary = NULL ;


/******************************************************************/
/* Exported function : Init                                       */
/******************************************************************/

IMPEXP void SpyDll_Init (DWORD dwServerProcessID, 
			 SPYDLL_NOTIFYPROC pfnNotifyProc) 
{
  TRACE_INFO (TEXT("Initializing DLL\n")) ;

  ASSERT (dwServerProcessID!=0) ;
  ASSERT (pfnNotifyProc!=NULL) ;

  if( sg_dwServerProcessID!=0 ) 
    {
      HANDLE	hProcess ;

      TRACE_WARNING (TEXT("Warning Spy DLL already initialized by process %d\n"),
		     sg_dwServerProcessID) ;
      
      hProcess = OpenProcess (0, FALSE, sg_dwServerProcessID) ;
      
      if( hProcess ) {
	TRACE_INFO (TEXT("The process %d is still running, I stop initialization.\n"), 
		    sg_dwServerProcessID) ;
	CloseHandle (hProcess) ;
	return ;
      }
      
      TRACE_INFO (TEXT("The process %d isn't running, I replace it\n"), sg_dwServerProcessID) ;
    }

  sg_dwServerProcessID	= dwServerProcessID ; 
  sg_pfnNotifyProc	= (PROC)pfnNotifyProc ;

  DbgPrint (TEXT("Server address = 0x%08X\n"), sg_pfnNotifyProc) ;
}


/******************************************************************/
/* Exported function : Uninit                                     */
/******************************************************************/

IMPEXP void SpyDll_Uninit ()
{
  TRACE_INFO (TEXT("Uninitializing DLL\n")) ;

  sg_dwServerProcessID	= 0 ;
  sg_pfnNotifyProc	= NULL ;
}


/******************************************************************/
/* Exported function : InjectInProcess                            */
/******************************************************************/

IMPEXP BOOL SpyDll_InjectInProcess (DWORD dwProcessId)
{
  BOOL		bResult = FALSE ;

  PROC pfnLoadLibrary = GetProcAddress (GetModuleHandle(TEXT("Kernel32.dll")), LOADLIBRARY_ANSI) ;

  ASSERT (dwProcessId!=GetCurrentProcessId()) ;
  ASSERT (dwProcessId!=sg_dwServerProcessID) ;

  TRACE_INFO (TEXT("Injecting process %u...\n"), dwProcessId) ;

  if( pfnLoadLibrary )
    {
      TCHAR szSpyDllPath[MAX_PATH] ;

      if( GetModuleFileName(g_hInstance, szSpyDllPath, MAX_PATH) )
	{
	  DWORD dwExitCode ;
	  int nSpyDllPathSize = (1 + _tcslen(szSpyDllPath)) * sizeof(TCHAR) ;
	  
	  if( SpyInt_CallRemoteFuncByPID (dwProcessId, pfnLoadLibrary,
					  szSpyDllPath, nSpyDllPathSize,
					  &dwExitCode) )
	    bResult = 0!=dwExitCode ;	      
	}
      else TRACE_ERROR (TEXT("GetModuleFileName failed\n")) ;
    }
  else TRACE_ERROR (TEXT("GetProcAddress failed\n")) ;

  if( bResult )
    TRACE_INFO (TEXT("Injection into %u succeed\n"), dwProcessId) ;
  else
    TRACE_INFO (TEXT("Injection into %u failed\n"), dwProcessId) ;
   
  return bResult ;
}


/******************************************************************/
/* Exported function : InjectInSuspendedThread                    */
/******************************************************************/

IMPEXP BOOL SpyDll_InjectInSuspendedThread (HANDLE hRemoteProcess, 
					    HANDLE hRemoteThread) 
{
  BOOL		bResult = FALSE ;

  BYTE pBootCode[] =
    {
      0x90,			// nop
      0x68,			// push
      0x00, 0x00, 0x00, 0x00,	// @"DllPath" 
      0x68,			// push
      0x00, 0x00, 0x00, 0x00,	// @EntryPoint
      0xFF, 0x25,		// jmp m32
      0x00, 0x00, 0x00, 0x00,	// @@LoadLibraryEx
      0x00, 0x00, 0x00, 0x00	// @LoadLibraryEx
    } ;

  PROC pfnLoadLibrary = g_pfnLoadLibrary ;
  //(PROC) GetProcAddress (GetModuleHandle(TEXT("Kernel32.dll")),LOADLIBRARY_ANSI) ;

  TRACE_INFO (TEXT("Injecting suspended process...\n")) ;

  if( pfnLoadLibrary )
    {
      TCHAR	szSpyDllPath[MAX_PATH] ;
      
      if( GetModuleFileName(g_hInstance, szSpyDllPath, MAX_PATH) )
	{
	  CONTEXT	context ;
	  context.ContextFlags = CONTEXT_FULL ;
    
	  if( GetThreadContext(hRemoteThread, &context) )
	    {	      
	      int	nSpyDllPathSize = (1 + _tcslen(szSpyDllPath)) * sizeof(TCHAR) ;
	      void	*pRemoteBuffer ;
	      int	nBootCodeSize = nSpyDllPathSize + sizeof(pBootCode) ; 

	      TRACE_INFO (TEXT("EIP=0x%08X\n"), context.Eip) ;
	      TRACE_INFO (TEXT("EAX=0x%08X\n"), context.Eax) ;	   
	      TRACE_INFO (TEXT("ESP=0x%08X\n"), context.Esp) ;
	      TRACE_INFO (TEXT("LoadLibrary=0x%08X\n"), pfnLoadLibrary) ;

	      pRemoteBuffer = VirtualAllocEx (hRemoteProcess,
					      NULL,
					      nBootCodeSize,
					      MEM_COMMIT,
					      PAGE_EXECUTE_READWRITE) ;

	      TRACE_INFO (TEXT("NEW=0x%08X\n"), pRemoteBuffer) ;
	      
	      if( pRemoteBuffer )
		{
		  DWORD		dwAddress ;

		  // add break point in boot code
#ifdef BREAK_IN_BOOT
		  pBootCode[0] = 0xCC ;				// int 3
#endif

		  // add dll path string address
		  dwAddress = (DWORD)(pRemoteBuffer+sizeof(pBootCode)) ; // @"DllPath"
		  memcpy (pBootCode+2, &dwAddress, 4) ;	

		  // add entry point address
		  dwAddress = (DWORD)(context.Eax) ;		// @EntryPoint
		  memcpy (pBootCode+7, &dwAddress, 4) ;	

		  // add @@LoadLibraryEx
		  dwAddress = (DWORD)(pRemoteBuffer+17) ;
		  memcpy (pBootCode+13, &dwAddress, 4) ;  

		  // add @LoadLibraryEx
		  dwAddress = (DWORD)pfnLoadLibrary ;
		  memcpy (pBootCode+17, &dwAddress, 4) ;  
	    
#ifdef SAVE_BOOT_INFO
		  {
		    FILE * fp = fopen ("C:\\boot.hex", "wb") ;
		    fwrite (pBootCode, sizeof(pBootCode), 1, fp) ;
		    fwrite (szSpyDllPath, nSpyDllPathSize, 1, fp) ;		    
		    fclose (fp) ;

		    TRACE_WARNING (TEXT("Boot code save to C:\\boot.hex\n")) ;
		  }
#endif

		  if( WriteProcessMemory (hRemoteProcess, pRemoteBuffer,
					  pBootCode, sizeof(pBootCode), NULL) 
		      && 
		      WriteProcessMemory (hRemoteProcess, pRemoteBuffer+sizeof(pBootCode),
					  szSpyDllPath, nSpyDllPathSize, NULL) 
		      )
		    {
		      context.Eax = (DWORD)pRemoteBuffer ;	  
		      
		      if( SetThreadContext (hRemoteThread, &context) )
			bResult = TRUE ;
		      else
			TRACE_ERROR (TEXT("SetThreadContext failed\n")) ;
		    }
		  else TRACE_ERROR (TEXT("WriteProcessMemory failed (error=%u)\n"), 
				    GetLastError()) ;
		}
	      else TRACE_ERROR (TEXT("VirtualAllocEx failed\n")) ;
	    }
	  else TRACE_ERROR (TEXT("GetThreadContext failed\n")) ;
	}
      else TRACE_ERROR (TEXT("GetModuleFileName failed\n")) ;
    }
  else TRACE_ERROR (TEXT("GetProcAddress failed\n")) ;

  if( bResult )
    TRACE_INFO (TEXT("Injection into suspended process succeed\n")) ;
  else
    TRACE_INFO (TEXT("Injection into suspended process failed\n")) ;

  return bResult ;
}


/******************************************************************/
/* Exported function : IsProcessInjected                          */
/******************************************************************/

IMPEXP HINSTANCE SpyDll_IsProcessInjected (DWORD dwProcessId,
					   LPCTSTR szDllFileName)
{
  TCHAR		szModulePath[MAX_PATH] ;
  DWORD		dwAccess = PROCESS_QUERY_INFORMATION|PROCESS_VM_READ ;
  HANDLE	hProcess ;
  HMODULE	*pModules ;
  int		i ;
  DWORD		dwModSize = 0 ;
  HINSTANCE	hSpyDllInstance = NULL ;

  ASSERT (dwProcessId!=GetCurrentProcessId()) ;
  ASSERT (dwProcessId!=sg_dwServerProcessID) ;
  
  hProcess = OpenProcess (dwAccess, FALSE, dwProcessId) ;
  
  if( ! hProcess ) {
    TRACE_ERROR (TEXT("OpenProcess failed (pid=%u, error=%d)\n"), dwProcessId, GetLastError()) ;
    return NULL ;
  }

  if( ! EnumProcessModules (hProcess, NULL, 0, &dwModSize) ) {
    TRACE_ERROR (TEXT("EnumProcessModules failed (pid=%u, error=%d)\n"), dwProcessId, GetLastError()) ;
    CloseHandle (hProcess) ;
    return NULL ;  
  }
  
  pModules = (HMODULE*) malloc (dwModSize) ;
  
  if( ! EnumProcessModules (hProcess, pModules, dwModSize, &dwModSize) ) {
    TRACE_ERROR (TEXT("EnumProcessModules failed (pid=%u, error=%d)\n"), dwProcessId, GetLastError()) ;
    free (pModules) ;
    CloseHandle (hProcess) ;
    return NULL ;  
  }

  for( i=0 ; i<(dwModSize/sizeof(HMODULE)) ; i++ )
    {
      GetModuleFileNameEx (hProcess, pModules[i], szModulePath, MAX_PATH) ;

      if( ! _tcsicmp(szModulePath,szDllFileName) )
	{
	  hSpyDllInstance = pModules[i] ;
	  break ;
	}
    }
  
  free (pModules) ;
  CloseHandle (hProcess) ;
  
  TRACE_INFO (TEXT("SpyDll %s in process %d\n"),
	      hSpyDllInstance ? TEXT("found") : TEXT("not found"), 
	      dwProcessId) ;
  
  return hSpyDllInstance ;
}



/******************************************************************/
/* Exported function : EjectFromProcess                           */
/******************************************************************/
/*
IMPEXP BOOL SpyDll_EjectFromProcess (DWORD dwProcessId, HINSTANCE hSpyDllInst)
{  
  BOOL		bResult = FALSE ;
  PROC		pfnFreeLibrary ;

  TRACE_INFO (TEXT("Un-injecting process %u...\n"), dwProcessId) ;

  pfnFreeLibrary = GetProcAddress (GetModuleHandle(TEXT("Kernel32.dll")), "FreeLibrary") ;

  if( pfnFreeLibrary )
    {
      DWORD dwExitCode ;
      
      if( SpyInt_CallRemoteFuncByPID (dwProcessId, pfnFreeLibrary,
				      hSpyDllInst, 0, &dwExitCode) )
	bResult = 0!=dwExitCode ; 
    }

  if( bResult )
    TRACE_INFO (TEXT("Un-injection from %u succeed\n"), dwProcessId) ;
  else
    TRACE_ERROR (TEXT("Un-injection from %u failed\n"), dwProcessId) ;

  return bResult ;
}
*/


DWORD WINAPI _SpyDll_SelfUnhookAndEject (void * pParam) 
{
  static	BOOL bRunning = FALSE ;
  HINSTANCE hSpyDllInst = (HINSTANCE) pParam ;
  
  TRACE ;

  if( bRunning ) return FALSE ;

  bRunning = TRUE ;

  Hooks_UninstallAll () ;

  FreeLibraryAndExitThread (hSpyDllInst,TRUE) ;

  return FALSE ;
}



IMPEXP BOOL SpyDll_EjectFromProcess (DWORD dwProcessId, HINSTANCE hRemoteSpyDll)
{  
  TCHAR		szDllPath[MAX_PATH] ;
  BOOL		bResult = FALSE ;
  PROC		pfnLocalEject, pfnRemoteEject ;
  HINSTANCE	hLocalSpyDll ;
  HANDLE	hRemoteProcess ;
  MODULEINFO	miRemote, miLocal ;
  BOOL		bSuccess ;
  DWORD		dwExitCode ;
  DWORD		dwDesiredAccess = PROCESS_CREATE_THREAD	| PROCESS_QUERY_INFORMATION | 
    PROCESS_VM_READ | PROCESS_VM_OPERATION | PROCESS_VM_WRITE ;

  if( dwProcessId==GetCurrentProcessId() )
    {
      TRACE_ERROR (TEXT("Can't eject from current process\n")) ;
      return FALSE ;
    }

  TRACE_INFO (TEXT("Un-injecting process %u...\n"), dwProcessId) ;

  hRemoteProcess = OpenProcess (dwDesiredAccess, FALSE, dwProcessId) ;
  
  if( hRemoteProcess )
    {
      bResult = 0!=GetModuleFileNameEx (hRemoteProcess, hRemoteSpyDll, szDllPath, MAX_PATH) ;

      if( bResult )
	{
	  TRACE_INFO (TEXT("Dll Path = %s\n"), szDllPath) ;

	  hLocalSpyDll = GetModuleHandle (szDllPath) ;

	  if( hLocalSpyDll )
	    {	      
	      bSuccess = GetModuleInformation (GetCurrentProcess(), hLocalSpyDll,  
					       &miLocal, sizeof(MODULEINFO)) ;
	      
	      if( bSuccess )
		{
		  pfnLocalEject = (PROC)_SpyDll_SelfUnhookAndEject ;
		  
		  bSuccess = GetModuleInformation (hRemoteProcess, hRemoteSpyDll,  
						   &miRemote, sizeof(MODULEINFO)) ;	  
		  
		  if( bSuccess )
		    {
		      pfnRemoteEject = (PROC)
			((DWORD)miRemote.lpBaseOfDll + (DWORD)pfnLocalEject - (DWORD)miLocal.lpBaseOfDll) ;
		      
		      if( SpyInt_CallRemoteFuncByHandle (hRemoteProcess, pfnRemoteEject,
							 hRemoteSpyDll, 0, &dwExitCode) )
			bResult = 0!=dwExitCode ; 
		    }
		  else TRACE_ERROR (TEXT("GetModuleInformation failed (error=%u)\n"),
				    GetLastError()) ;
		}
	      else TRACE_ERROR (TEXT("GetModuleInformation failed (error=%u)\n"),
				GetLastError()) ;
	    }
	  else TRACE_ERROR (TEXT("GetModuleHandle failed (error=%u)\n"),
				GetLastError()) ;
	}    
      else TRACE_ERROR (TEXT("GetModuleFileNameEx failed (error=%u)\n"),
			GetLastError()) ;
      
      CloseHandle (hRemoteProcess) ;
    }
  else TRACE_ERROR (TEXT("OpenProcess failed (pid=%u, error=%u)\n"), 
		    dwProcessId, GetLastError()) ; 
    
  if( bResult )
    TRACE_INFO (TEXT("Un-injection from %u succeed\n"), dwProcessId) ;
  else
    TRACE_ERROR (TEXT("Un-injection from %u failed\n"), dwProcessId) ;

  return bResult ;
}



/******************************************************************/
/* Exported function : DllMain                                    */
/******************************************************************/

BOOL WINAPI DllMain (HINSTANCE hInstance, DWORD dwReason, VOID * pReserved)
{
  switch( dwReason )
    {
    case DLL_PROCESS_ATTACH:
      {
	HANDLE		hProcess = GetCurrentProcess() ;

	DisableThreadLibraryCalls (hInstance) ;
      
	g_hInstance = hInstance ;
	g_hModule = GetModuleHandle (NULL) ;
	g_pfnLoadLibrary = (PROC) GetProcAddress (GetModuleHandle(TEXT("Kernel32.dll")), 
						  LOADLIBRARY_ANSI) ;
	

	GetModuleInformation (hProcess, hInstance,
			      &g_miInstance, sizeof(MODULEINFO)) ;
	
	if( ! Hooks_Init () ) {
	  TRACE_ERROR (TEXT("Hooks_Init failed\n")) ;
	  return FALSE ;
	}
	Hooks_InstallIntoMappedModules () ;

	DbgPrint (TEXT("Hooks installed\n")) ;
	
	TRACE_INFO(TEXT("0x%08X - Spy_CreateProcessW\n"), (DWORD)Spy_CreateProcessW) ;
	TRACE_INFO(TEXT("0x%08X - SpyInt_CatchCall\n"), (DWORD)SpyInt_CatchCall) ;
	TRACE_INFO(TEXT("0x%08X - SpyInt_NotifyServer\n"), (DWORD)SpyInt_NotifyServer) ;
	TRACE_INFO(TEXT("0x%08X - SpyInt_CallRemoteFuncByPID\n"), (DWORD)SpyInt_CallRemoteFuncByPID) ;
	TRACE_INFO(TEXT("0x%08X - SpyInt_CallRemoteFuncByHandle\n"), (DWORD)SpyInt_CallRemoteFuncByHandle) ;
      }
      return TRUE ;

      
    case DLL_PROCESS_DETACH:
      {
	DbgPrint (TEXT("Removing hooks...\n")) ;
	
	Hooks_UninstallAll () ;
	
	Hooks_Uninit () ;
	
	DbgPrint (TEXT("Hooks removed\n")) ;
      }
      return TRUE ; 
    }

  return TRUE ;
}
